"""


Version: 0.1
Author: lk
Date: 2022-03-08 18:35
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn.pytorch.conv import GINConv
from dgl.nn.pytorch.glob import AvgPooling, MaxPooling, SumPooling

class SElayer(nn.Module):
    def __init__(self,in_channels,se_channels):
        super().__init__()
        self.encoder_decoder=nn.Sequential(
            nn.Linear(in_channels,se_channels),
            nn.ELU(),
            nn.Linear(se_channels,in_channels),
            nn.Sigmoid()
        )
    def forward(self,x):
        x_global = torch.mean(x,dim=0)
        s=self.encoder_decoder(x_global)
        return s * x




if __name__ == '__main__':
    x=torch.randn(5,12)
    in_channels=12
    se_channels=5
    se = SElayer(in_channels,se_channels)
    print(x)
    print(se(x))